﻿// test_chacha.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#define _CRT_SECURE_NO_WARNINGS

#include <stdio.h>
#include <inttypes.h>
#include <iostream>
#include <vector>
#include "chacha.h"
#include "rand.h"

using namespace std;


__device__ uint32_t rotate_right(uint32_t number, uint32_t shiftAmount)
{
	return (number >> shiftAmount) | (number << (32 - shiftAmount));
}

__device__ void expand_seed(uint64_t state, uint8_t seed[32])
{
	// We use PCG32 to generate a u32 sequence, and copy to the seed
	auto pcg32 = [](uint64_t& state, uint8_t* ret) {
		const uint64_t MUL = 6364136223846793005ULL;
		const uint64_t INC = 11634580027462260723ULL;

		// We advance the state first (to get away from the input value,
		// in case it has low Hamming Weight).
		state = state * MUL + INC;

		// Use PCG output function with to_le to generate x:
		uint32_t xorshifted = uint32_t(((state >> 18) ^ state) >> 27);
		uint32_t rot = uint32_t(state >> 59);
		
		uint32_t x = rotate_right(xorshifted, rot);
		//to_le_bytes()
		memcpy(ret, &x, 4);
	};

	for (int i = 0; i < 32; i += 4)
	{
		pcg32(state, seed + i);
	}
}

//MyRandom
__device__ MyRandom::MyRandom()
{
	results.resize(2048 / 32);
	index = (int)results.size();
	m_log = false;
	m_genTime = 0;
}

__device__ void MyRandom::setDebug(bool flag)
{
	m_log = flag;
}

__device__ int MyRandom::genTime()
{
	return m_genTime;
}

__device__ uint64_t MyRandom::init(CuByteArray in_epoch_hash, CuByteArray in_address, uint64_t counter)
{
	const int addr_size = 32;
	const int hash_size = 32;
	assert(in_address.size() == addr_size && hash_size == in_epoch_hash.size());
	uint8_t hash_input[8 + addr_size + 8] = {};

	memcpy(hash_input, in_epoch_hash.data(), 8);
	memcpy(hash_input + 8, in_address.data(), in_address.size());
	memcpy(hash_input + 8 + addr_size, &counter, 8);

	uint64_t solution_id = sha256_u64(hash_input, sizeof(hash_input));
	uint8_t seed[32] = {};
	expand_seed(solution_id, seed);

	chacha20_init(&m_rng, (const uint32_t*)seed, 0);
	return solution_id;
}

__device__ uint64_t MyRandom::sha256_u64(uint8_t* buffer, int len)
{
	/*
	QByteArrayView input(buffer, len);
	QByteArray hash = QCryptographicHash::hash(input, QCryptographicHash::Sha256);
	hash = QCryptographicHash::hash(hash, QCryptographicHash::Sha256);

	uint64_t* pu64 = (uint64_t*)hash.data();
	return *pu64;
	*/
	return 0;
}

__device__ Boolean MyRandom::rand_bool()
{
	return Boolean::create(Public,int32_t(this->next_u32()) < 0);
}

__device__ Integer MyRandom::rand_i8()
{
	return Integer::create(Type::i8, Public, (int8_t)this->next_u32());
}

__device__ Integer MyRandom::rand_i16()
{
	return Integer::create(Type::i16, Public, (int16_t)this->next_u32());
}

__device__ Integer MyRandom::rand_i32()
{
	return Integer::create(Type::i32, Public, (int32_t)next_u32());
}

__device__ Integer MyRandom::rand_i64()
{
	return Integer::create(Type::i64, Public, (int64_t)next_u64());
}

__device__ Integer MyRandom::rand_i128()
{
	uint64_t buffer[2] = {};
	buffer[0] = this->next_u64();
	buffer[1] = this->next_u64();
	
	char* p_buffer = (char*)&buffer;
	CuByteArray bytes;
	for (int i = 0; i < 16; i++)
		bytes.push_back(p_buffer[i]);

	auto bits = bytes_to_bits_le(bytes);
	bits[0].mode = Public;
	return Integer::fromBitsLe(Type::i128, bits);
}

__device__ Integer MyRandom::rand_u8()
{
	return Integer::create(Type::u8, Public, (uint8_t)this->next_u32());
}

__device__ Integer MyRandom::rand_u16()
{
	return Integer::create(Type::u16, Public, (uint16_t)this->next_u32());
}

__device__ Integer MyRandom::rand_u32()
{
	return Integer::create(Type::u32, Public, this->next_u32());
}

__device__ Integer MyRandom::rand_u64()
{
	return Integer::create(Type::u64, Public, this->next_u64());
}

//D:\work\snarkVM\fields\src\macros.rs impl_primefield_standard_sample
__device__ Field MyRandom::rand_field()
{
	auto REPR_SHAVE_BITS = 3;
	uint64_t var[4] = {};
	uint256_t b;

	while (true)
	{
		for (int i = 0; i < 4; i++) {
			var[i] = this->next_u64();
		}
		var[3] &= UINT64_MAX >> REPR_SHAVE_BITS;

		
		if (b < *P::MODULUS)
			break;
	}

	b = to_bigint(b);
	return Field::create(Mode::Public,b);
}

__device__ uint32_t MyRandom::next_u32() {
	if (this->index >= this->results.size()) {
		this->generate_and_set(0);
	}

	auto value = this->results[this->index];
	this->index += 1;
	return value;
}

__device__ uint64_t MyRandom::next_u64()
{
	auto read_u64 = [](uint32_t* results, int index) -> uint64_t {
		uint32_t* data = results + index;
		return uint64_t(data[1]) << 32 | uint64_t(data[0]);
		};

	int len = (int)this->results.size();
	int index = this->index;
	if (index < len - 1) {
		this->index += 2;
		// Read an u64 from the current index
		return read_u64(this->results.data(), index);
	}
	else if (index >= len)
	{
		this->generate_and_set(2);
		return read_u64(this->results.data(), 0);
	}
	else
	{
		uint64_t x = this->results[len - 1];
		this->generate_and_set(1);
		uint64_t y = this->results[0];
		return (y << 32) | x;
	}
}

__device__ void MyRandom::generate_and_set(int index)
{
	for (int i = 0; i < results.size(); i++)
		results[i] = 0;

	uint8_t* array = (uint8_t*)results.data();
	for (int i = 0; i < 4; i++)
	{
		chacha_fill_u8(&m_rng, array + i * 64, 64);
	}
	this->index = index;
	m_genTime++;
}

__device__ uint64_t mac_with_carry(uint64_t a, uint64_t b, uint64_t c, uint64_t* carry)
{
	auto tmp = (uint128_t(a)) + uint128_t(b) * uint128_t(c) + uint128_t(*carry);
	uint64_t *p_carry = (uint64_t *)(tmp.data + 2);
	*carry = *p_carry;
	return (uint64_t)tmp;
}

__device__ uint256_t MyRandom::to_bigint(uint256_t in)
{
	uint64_t r[4] = {};
	uint64_t P_MODULUS[4] = {};
	memcpy(r,in.data,sizeof(r));
	memcpy(P_MODULUS,P::MODULUS->data,sizeof(P_MODULUS));

	// Montgomery Reduction
	uint64_t k = r[0] * (P::INV);
	uint64_t carry = 0;
	mac_with_carry(r[0], k, P_MODULUS[0], &carry);
	r[1] = mac_with_carry(r[1], k, P_MODULUS[1], &carry);
	r[2] = mac_with_carry(r[2], k, P_MODULUS[2], &carry);
	r[3] = mac_with_carry(r[3], k, P_MODULUS[3], &carry);
	r[0] = carry;

	k = r[1] * (P::INV);
	carry = 0;
	mac_with_carry(r[1], k, P_MODULUS[0], &carry);
	r[2] = mac_with_carry(r[2], k, P_MODULUS[1], &carry);
	r[3] = mac_with_carry(r[3], k, P_MODULUS[2], &carry);
	r[0] = mac_with_carry(r[0], k, P_MODULUS[3], &carry);
	r[1] = carry;

	k = r[2] * (P::INV);
	carry = 0;
	mac_with_carry(r[2], k, P_MODULUS[0], &carry);
	r[3] = mac_with_carry(r[3], k, P_MODULUS[1], &carry);
	r[0] = mac_with_carry(r[0], k, P_MODULUS[2], &carry);
	r[1] = mac_with_carry(r[1], k, P_MODULUS[3], &carry);
	r[2] = carry;

	k = r[3] * (P::INV);
	carry = 0;
	mac_with_carry(r[3], k, P_MODULUS[0], &carry);
	r[0] = mac_with_carry(r[0], k, P_MODULUS[1], &carry);
	r[1] = mac_with_carry(r[1], k, P_MODULUS[2], &carry);
	r[2] = mac_with_carry(r[2], k, P_MODULUS[3], &carry);
	r[3] = carry;

	uint256_t ret;
	return ret;
}